from datetime import datetime, date
import backtrader as bt
import akshare as ak
import pandas as pd
import numpy as np
import quantstats as qs
import pyfolio as pf
import webbrowser
from sqlalchemy import create_engine, distinct, or_, and_
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
from sklearn.linear_model import LinearRegression
import warnings

warnings.filterwarnings('ignore')

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def get_stocks():
    sql = "SELECT STK_CODE FROM ETF_INDEX_INFO"
    r = exec_sql(sql)
    df1 = pd.DataFrame(r, columns=['代码'])
    df2 = ak.fund_etf_spot_em()
    res = pd.merge(df1, df2, on=['代码'])
    res = res.loc[res['最新价'] < 10, :]
    stocks = res['代码'].to_list()
    # stocks = [tok for tok in stocks if tok != '512870']

    return stocks


def get_data(stocks, freq='day', source='db'):
    if freq == '1min':
        table_name = 'STOCK_MIN_DATA'
    elif freq == '15min':
        table_name = 'STOCK_15MIN_DATA'
    else:
        table_name = 'ETF_DAY_DATA'  # HIS_DAILY_DATA

    res = dict()
    date_data = set_date()
    st = date_data['st']
    et = date_data['et']
    st_d = date_data['st_d']
    et_d = date_data['et_d']

    if source == 'db':
        print('db')
        stocks_str = ','.join(f"'{tok}'" for tok in stocks)
        sql = f"SELECT code, datetime, [close], [open], low, high, volume FROM {table_name} WHERE code in ({stocks_str})"
        r = exec_sql(sql)
        df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume'])
        # df['datetime'] = df['datetime'].map(time_map)

        for code in stocks:
            if code == '399300':
                sig_df = ak.stock_zh_index_daily(symbol="sz399300")
                sig_df = sig_df.loc[(sig_df.date >= st_d) & (sig_df.date <= et_d), :]
                sig_df.rename(columns={'date': 'datetime'}, inplace=True)
                sig_df.insert(0, 'code', code)
                sig_df = sig_df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
                # sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            else:
                sig_df = df.loc[df['code'] == code]
            sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', drop=False, inplace=True)
            if len(sig_df) <= 220:
                continue
            print(f'code={code}->len={len(sig_df)}')
            res[code] = sig_df
    else:
        print('internet')
        for code in stocks:
            if code == '399300':
                df = ak.stock_zh_index_daily(symbol="sz399300")
                df = df.loc[(df.date >= st_d) & (df.date <= et_d), :]
                df.rename(columns={'date': 'datetime'}, inplace=True)
            else:
                df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
                df.rename(columns={'日期': 'datetime', '开盘': 'open', '收盘': 'close', '最高': 'high',
                                   '最低': 'low', '成交量': 'volume'}, inplace=True)

            df.insert(0, 'code', code)
            df = df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
            df['datetime'] = pd.to_datetime(df['datetime'])
            df.set_index('datetime', inplace=True)
            print(f'code={code}->len={len(df)}')
            res[code] = df

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = datetime.strptime(str(x), '%Y-%m-%d')
    # t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


stocks_info = dict()
init_cash = 0.5e6


def run_backtrade(stocks):
    cerebro = bt.Cerebro()
    cerebro.addstrategy(GridReversionStrategy)
    data_all = get_data(stocks, source='db')
    index = stocks[0]
    data0 = data_all[index]
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    date_data = set_date()
    st_dt = date_data['st_dt']
    et_dt = date_data['et_dt']

    for i, code in enumerate(data_all):
        if code == index:
            df = data_all[code]
            df = df.loc[:, ['close']]
            df.rename(columns={'close': '399300'}, inplace=True)
            continue
        data = data_all[code]
        print(f"i={i}->len={len(data)}->{code} Done !")
        data_bf = data.loc[data['datetime'] < st_dt]

        if len(data_bf) > 0:
            info = data_bf.loc[data_bf['high'].idxmax(), :]
            stocks_info[code] = {'hdt': info['datetime'], 'high': info['high']}
        datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_dt, todate=et_dt)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    # cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    # st = strat.params.st
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    # print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    print(st_dt.date())
    returns = returns[returns.index.date >= st_dt.date()]
    returns.index = returns.index.date
    returns.index = pd.to_datetime(returns.index)

    print(f'returns={returns}')
    html_file = 'test.html'
    qs.reports.html(returns, benchmark=df, title='grid', output=html_file, download_filename=html_file)
    qs.reports.basic(returns)
    f = webbrowser.open(html_file)
    # pf.create_full_tear_sheet(returns)
    # cerebro.plot()


class GridReversionStrategy(bt.Strategy):
    params = (
        ('period', 25),
        ('init_money', 3e4),
        ('add_money', 1e4),
        ('buy_ratio', -0.35),
        ('add_ratio', -0.1),
        ('sell_ratio', 0.2),
        ('printlog', False),
        ('maperiod', 26),
        ('times', 1),
        ('span', 220),
        ('st', datetime(2014, 1, 1)),
    )

    def __init__(self):
        self.order = None
        self.daily_pnl = dict()
        self.trade_info = list()
        self.pos_info = dict()
        self.trail_dict = dict()
        self.buy_first = dict()  # 每支ETF是否开始第一天就跌破之前最高价的buy_ratio
        self.cache_log = list()
        self.TR = dict()
        self.ATR = dict()
        # print(f'stocks_info={stocks_info}')

        stocks = list()
        for data in self.datas:
            code = data._name
            if code != '399300':
                stocks.append(code)
                if code in stocks_info:
                    self.trail_dict[code] = {'hdt': stocks_info[code]['hdt'], 'high': stocks_info[code]['high'],
                                             'ldt': 0, 'low': 0, 'cdt': 0, 'cost': 0, 'buy_count': 0, 'trend': 'x',
                                             'pos': dict(), 'grid': dict(), 'counter': 0, 'trail_high': False}
                else:
                    self.trail_dict[code] = {'hdt': 0, 'high': 0, 'ldt': 0, 'low': 0, 'cdt': 0, 'cost': 0,
                                             'buy_count': 0, 'trend': 'x', 'pos': dict(),
                                             'grid': dict(), 'counter': 0, 'trail_high': False}
                self.TR[code] = bt.indicators.Max((data.high - data.low), abs(data.close(-1) - data.high),
                                                        abs(data.close(-1) - data.low))

                self.ATR[code] = bt.indicators.SMA(self.TR[code], period=self.p.maperiod, subplot=False)
                self.buy_first[code] = True

    def downcast(self, price, style='add'):
        res = self.p.add_money / price // 100 * 100 if style == 'add' else self.p.init_money / price // 100 * 100
        return res

    def set_grid(self, code, price):
        self.trail_dict[code]['grid']['lev0'] = price
        self.trail_dict[code]['grid']['lev_m'] = price * (1 + 2.7 * self.p.add_ratio)
        self.trail_dict[code]['grid']['lev1'] = price * (1 + self.p.buy_ratio)
        self.trail_dict[code]['grid']['lev2'] = self.trail_dict[code]['grid']['lev1'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev3'] = self.trail_dict[code]['grid']['lev2'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev4'] = self.trail_dict[code]['grid']['lev3'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev5'] = self.trail_dict[code]['grid']['lev4'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev6'] = self.trail_dict[code]['grid']['lev5'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev7'] = self.trail_dict[code]['grid']['lev6'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev8'] = self.trail_dict[code]['grid']['lev7'] * (1 + self.p.add_ratio)
        self.log(f"high={price}->lev_m={price * (1 + self.p.add_ratio)}->lev1={price * (1 + self.p.buy_ratio)}->"
                 f"lev2={self.trail_dict[code]['grid']['lev1'] * (1 + self.p.add_ratio)}->"
                 f"lev3={self.trail_dict[code]['grid']['lev2'] * (1 + self.p.add_ratio)}->"
                 f"lev4={self.trail_dict[code]['grid']['lev3'] * (1 + self.p.add_ratio)}")
        self.log(50 * '-')

    def prenext(self):
        self.next()

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            content = f'{dt.isoformat()}, {txt}'
            self.cache_log.append(content)
            print(content)

    @staticmethod
    def percent(today, yesterday):
        return float(today - yesterday) / today

    def next(self):
        if self.order:
            return

        cur_dt = self.datas[0].datetime.datetime(0)
        dt_int = int(cur_dt.strftime('%Y%m%d'))

        if dt_int not in self.daily_pnl:
            self.daily_pnl[dt_int] = round(self.broker.getvalue() - init_cash, 2)

        for data in self.datas:
            code = data._name
            try:
                dt = data.datetime.datetime(0)  # 获取当前的回测时间点
            except:
                print(f'{code}-{dt}')
            # dt != cur_dt 防止停牌日触发买卖
            if code == '399300' or dt != cur_dt:
                continue
            elif not self.trail_dict[code]['high'] and self.trail_dict[code]['counter'] < self.p.span:
                self.trail_dict[code]['counter'] = self.trail_dict[code]['counter'] + 1
                continue

            if not self.buy_first[code] and data.high[0] > self.trail_dict[code]['high']:
                self.trail_dict[code]['hdt'] = dt
                self.trail_dict[code]['high'] = data.high[0]
                self.set_grid(code, data.high[0])
                self.trail_dict[code]['trail_high'] = True
                continue
            elif self.buy_first[code]:
                if self.trail_dict[code]['high']:
                    self.set_grid(code, self.trail_dict[code]['high'])
                else:
                    high_list = data.lines.high.get(ago=-1, size=self.trail_dict[code]['counter'] - 1)
                    self.trail_dict[code]['high'] = max(high_list)
                    self.set_grid(code, self.trail_dict[code]['high'])

            if self.buy_first[code]:
                available_cash = self.broker.getcash()  # 获取当前可用资金
                available_amount = available_cash * 0.95 / data.low[0] // 100 * 100
                amount = 0
                buy_count = 0
                self.log(f'buy_first={self.buy_first[code]}->available_amount={available_amount}')
                if data.low[0] < self.trail_dict[code]['grid']['lev8']:
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount4 = 6 * amount2
                    amount5 = 1.1 * amount4
                    amount6 = 1.1 * amount5
                    amount7 = 1.1 * amount6
                    amount8 = 1.1 * amount7
                    if available_amount >= amount1 + amount2 + amount3 + amount4 + amount5 + amount6 + amount7 + amount8:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5 + amount6 + amount7 + amount8
                        buy_count = 8
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.trail_dict[code]['pos'][6] = amount6
                        self.trail_dict[code]['pos'][7] = amount7
                        self.trail_dict[code]['pos'][8] = amount8
                        self.log(f'buy lev8->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}->a6={amount6}->a7={amount7}->a8={amount8}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4 + amount5 + amount6 + amount7:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5 + amount6 + amount7
                        buy_count = 7
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.trail_dict[code]['pos'][6] = amount6
                        self.trail_dict[code]['pos'][7] = amount7
                        self.log(f'buy lev8->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}->a6={amount6}->a7={amount7}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4 + amount5 + amount6:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5 + amount6
                        buy_count = 6
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.trail_dict[code]['pos'][6] = amount6
                        self.log(f'buy lev8->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}->a6={amount6}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4 + amount5:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5
                        buy_count = 5
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.log(f'buy lev8->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4:
                        amount = amount1 + amount2 + amount3 + amount4
                        buy_count = 4
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.log(f'buy lev8->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}')
                    elif available_amount >= amount1 + amount2 + amount3:
                        amount = amount1 + amount2 + amount3
                        buy_count = 3
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.log(f'buy lev8->a1={amount1}->a2={amount2}->a3={amount3}')
                    elif available_amount >= amount1 + amount2:
                        amount = amount1 + amount2
                        buy_count = 2
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.log(f'buy lev8>a1={amount1}->a2={amount2}')
                    elif available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev8->a1_1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev8->a1_2={amount1}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev7']:
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount4 = 6 * amount2
                    amount5 = 1.1 * amount4
                    amount6 = 1.1 * amount5
                    amount7 = 1.1 * amount6
                    if available_amount >= amount1 + amount2 + amount3 + amount4 + amount5 + amount6 + amount7:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5 + amount6 + amount7
                        buy_count = 7
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.trail_dict[code]['pos'][6] = amount6
                        self.trail_dict[code]['pos'][7] = amount7
                        self.log(f'buy lev7->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}->a6={amount6}->a7={amount7}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4 + amount5 + amount6:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5 + amount6
                        buy_count = 6
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.trail_dict[code]['pos'][6] = amount6
                        self.log(f'buy lev7->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}->a6={amount6}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4 + amount5:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5
                        buy_count = 5
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.log(f'buy lev7->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4:
                        amount = amount1 + amount2 + amount3 + amount4
                        buy_count = 4
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.log(f'buy lev7->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}')
                    elif available_amount >= amount1 + amount2 + amount3:
                        amount = amount1 + amount2 + amount3
                        buy_count = 3
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.log(f'buy lev7->a1={amount1}->a2={amount2}->a3={amount3}')
                    elif available_amount >= amount1 + amount2:
                        amount = amount1 + amount2
                        buy_count = 2
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.log(f'buy lev7->a1={amount1}->a2={amount2}')
                    elif available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev7->a1_1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev7->a1_2={amount1}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev6']:
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount4 = 6 * amount2
                    amount5 = 1.1 * amount4
                    amount6 = 1.1 * amount5
                    if available_amount >= amount1 + amount2 + amount3 + amount4 + amount5 + amount6:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5 + amount6
                        buy_count = 6
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.trail_dict[code]['pos'][6] = amount6
                        self.log(f'buy lev6->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}->a6={amount6}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4 + amount5:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5
                        buy_count = 5
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.log(f'buy lev6->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4:
                        amount = amount1 + amount2 + amount3 + amount4
                        buy_count = 4
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.log(f'buy lev6->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}')
                    elif available_amount >= amount1 + amount2 + amount3:
                        amount = amount1 + amount2 + amount3
                        buy_count = 3
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.log(f'buy lev6->a1={amount1}->a2={amount2}->a3={amount3}')
                    elif available_amount >= amount1 + amount2:
                        amount = amount1 + amount2
                        buy_count = 2
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.log(f'buy lev6->a1={amount1}->a2={amount2}')
                    elif available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev6->a1_1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev6->a1_2={amount1}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev5']:
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount4 = 6 * amount2
                    amount5 = 1.1 * amount4
                    if available_amount >= amount1 + amount2 + amount3 + amount4 + amount5:
                        amount = amount1 + amount2 + amount3 + amount4 + amount5
                        buy_count = 5
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.trail_dict[code]['pos'][5] = amount5
                        self.log(f'buy lev5->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}'
                                 f'->a5={amount5}')
                    elif available_amount >= amount1 + amount2 + amount3 + amount4:
                        amount = amount1 + amount2 + amount3 + amount4
                        buy_count = 4
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.trail_dict[code]['pos'][4] = amount4
                        self.log(f'buy lev5->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}')
                    elif available_amount >= amount1 + amount2 + amount3:
                        amount = amount1 + amount2 + amount3
                        buy_count = 3
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.log(f'buy lev5->a1={amount1}->a2={amount2}->a3={amount3}')
                    elif available_amount >= amount1 + amount2:
                        amount = amount1 + amount2
                        buy_count = 2
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.log(f'buy lev5->a1={amount1}->a2={amount2}')
                    elif available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev5->a1_1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev5->a1_2={amount1}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev4']:
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount4 = 6 * amount2
                    if available_amount >= amount1 + amount2 + amount3 + amount4:
                       amount = amount1 + amount2 + amount3 + amount4
                       buy_count = 4
                       self.trail_dict[code]['pos'][1] = amount1
                       self.trail_dict[code]['pos'][2] = amount2
                       self.trail_dict[code]['pos'][3] = amount3
                       self.trail_dict[code]['pos'][4] = amount4
                       self.log(f'buy lev4->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}')
                    elif available_amount >= amount1 + amount2 + amount3:
                        amount = amount1 + amount2 + amount3
                        buy_count = 3
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.log(f'buy lev4->a1={amount1}->a2={amount2}->a3={amount3}')
                    elif available_amount >= amount1 + amount2:
                        amount = amount1 + amount2
                        buy_count = 2
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.log(f'buy lev4->a1={amount1}->a2={amount2}')
                    elif available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev4->a1_1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev4->a1_2={amount1}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev3']:
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount = 0
                    buy_count = 0
                    if available_amount >= amount1 + amount2 + amount3:
                        amount = amount1 + amount2 + amount3
                        buy_count = 3
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.trail_dict[code]['pos'][3] = amount3
                        self.log(f'buy lev3->a1={amount1}->a2={amount2}->a3={amount3}')
                    elif available_amount >= amount1 + amount2:
                        amount = amount1 + amount2
                        buy_count = 2
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.log(f'buy lev3->a1={amount1}->a2={amount2}')
                    elif available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev3->a1_1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                elif data.low[0] < self.trail_dict[code]['grid']['lev2']:
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount = 0
                    buy_count = 0
                    if available_amount >= amount1 + amount2:
                        amount = amount1 + amount2
                        buy_count = 2
                        self.trail_dict[code]['pos'][1] = amount1
                        self.trail_dict[code]['pos'][2] = amount2
                        self.log(f'buy lev2->a1={amount1}->a2={amount2}')
                    elif available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev2->a1_1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                elif data.low[0] < self.trail_dict[code]['grid']['lev1']:
                    amount1 = self.downcast(data.low[0])
                    amount = 0
                    buy_count = 0
                    self.log(f'all_amount={amount1}')
                    if available_amount >= amount1:
                        amount = amount1
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev1->a1={amount1}')
                    elif available_amount >= 0:
                        amount = available_amount
                        buy_count = 1
                        self.trail_dict[code]['pos'][1] = amount
                        self.log(f'buy lev1->a1={amount1}')
                else:
                    self.buy_first[code] = False

                if amount > 0:
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.trail_dict[code]['buy_count'] = buy_count
            else:
                if self.trail_dict[code]['buy_count'] == 0 and data.low[0] < self.trail_dict[code]['grid']['lev1']:
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    if self.trail_dict[code]['trail_high']:
                        self.trail_dict[code]['trail_high'] = False
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev1->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 1 and data.low[0] < self.trail_dict[code]['grid']['lev2']:
                    amount = self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev2->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 2 and data.low[0] < self.trail_dict[code]['grid']['lev3']:
                    amount = 3 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev3->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 3 and data.low[0] < self.trail_dict[code]['grid']['lev4']:
                    amount = 6 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev4->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 4 and data.low[0] < self.trail_dict[code]['grid']['lev5']:
                    amount = 6 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev5->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 5 and data.low[0] < self.trail_dict[code]['grid']['lev6']:
                    amount = 6 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev6->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 6 and data.low[0] < self.trail_dict[code]['grid']['lev7']:
                    amount = 6 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev7->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 7 and data.low[0] < self.trail_dict[code]['grid']['lev8']:
                    amount = 6 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev8->xx_a={amount}')

                # 卖出
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 8 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev6']:
                    amount = self.trail_dict[code]['pos'][8]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 7 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev5']:
                    amount = self.trail_dict[code]['pos'][7]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 6 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev4']:
                    amount = self.trail_dict[code]['pos'][6]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 5 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev3']:
                    amount = self.trail_dict[code]['pos'][5]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 4 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev2']:
                    amount = self.trail_dict[code]['pos'][4]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 3 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev1']:
                    amount = self.trail_dict[code]['pos'][3]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_2'
                    self.log(f'sell lev1->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 2 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev_m']:
                    amount = self.trail_dict[code]['pos'][2]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_3'
                    self.log(f'sell lev_m->in_a={amount}')
                # elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and \
                #         data.high[
                #             0] > self.trail_dict[code]['grid']['lev0'] * 0.85:
                #     amount = self.trail_dict[code]['pos'][1]
                #     self.order = self.sell(data=data, size=amount, name=code)
                #     self.trail_dict[code]['ldt'] = dt
                #     self.trail_dict[code]['low'] = data.high[0]
                #     self.trail_dict[code]['trend'] = 'up1_4'
                #     self.log(f'sell lev0->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and \
                        data.close[0] < self.ATR[code] * self.p.times:
                    amount = self.trail_dict[code]['pos'][1]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_4'
                    high_list = data.lines.high.get(ago=-1, size=self.p.span)
                    self.trail_dict[code]['high'] = max(high_list)
                    self.set_grid(code, self.trail_dict[code]['high'])
                    self.log(f'sell lev0->in_a={amount}')
                elif self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and \
                        data.close[0] < self.trail_dict[code]['high'] * 0.9:
                    self.trail_dict[code]['trail_high'] = False
                    self.order = self.close(data=data, name=code)
                    high_list = data.lines.high.get(ago=-1, size=self.p.span)
                    self.trail_dict[code]['high'] = max(high_list)
                    self.set_grid(code, self.trail_dict[code]['high'])
                    self.log('set_false')
                # elif self.trail_dict[code]['trail_high'] and data.close[0] < self.ATR[code] * self.p.times:
                #     self.trail_dict[code]['trail_high'] = False
                #     self.order = self.close(data=data, name=code)
                #     self.log('set_false')

                if data.low[0] < self.trail_dict[code]['low']:
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]

                # self.log(f"trail_high={self.trail_dict[code]['trail_high'] }")

    def notify_order(self, order):
        code = order.data._name

        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:

            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={code} BUY EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f'Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')

                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'BUY', round(order.executed.price, 2), order.executed.size,
                              round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                if self.buy_first[code]:
                    self.buy_first[code] = False
                else:
                    self.trail_dict[code]['buy_count'] += 1
                    self.trail_dict[code]['pos'][self.trail_dict[code]['buy_count']] = order.executed.size
                self.trail_dict[code]['cdt'] = int(bt.num2date(order.executed.dt).strftime('%Y%m%d'))
                self.trail_dict[code]['cost'] = round(order.executed.price, 2)
            else:
                self.log(
                    f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} SELL EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'SELL', round(order.executed.price, 2),
                              order.executed.size, round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                self.trail_dict[code]['cdt'] = int(bt.num2date(order.executed.dt).strftime('%Y%m%d'))
                self.trail_dict[code]['cost'] = round(order.executed.price, 2)
                del self.trail_dict[code]['pos'][self.trail_dict[code]['buy_count']]
                self.trail_dict[code]['buy_count'] -= 1
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f'Order Canceled/Margin/Rejected->code={code}')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {trade.pnl:.2f}, NET {trade.pnlcomm:.2f}')

    def stop(self):
        print('Strategy Finish!!!')

        # with open(r'C:\Users\AndrewX\Desktop\cache_log.txt', 'w') as f:
        #     res = '\n'.join(self.cache_log)
        #     f.write(res)
        doprint = (False, True)[1]

        if doprint:
            daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
            append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
                                     columns=['DATE_T', 'PNL'])
            daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
            trade_df = pd.DataFrame(self.trade_info,
                                    columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
            pnl_table = 'TEST_DAILY_PNL'
            trade_table = 'TEST_TRADE_INFO'
            # daily_pnl_df = daily_pnl_df.loc[(daily_pnl_df['DATE_T'] >= 20230428) | (daily_pnl_df['DATE_T'] == 101)]
            insert_db(daily_pnl_df, pnl_table)
            insert_db(trade_df, trade_table)
            print(self.trail_dict)
            print('insert table ok')


def set_date():
    st = "20140101"
    et = "20230816"
    st_dt = datetime.strptime(st, '%Y%m%d')
    et_dt = datetime.strptime(et, '%Y%m%d')
    st_d = st_dt.date()
    et_d = et_dt.date()
    res = {'st': st, 'et': et, 'st_d': st_d, 'et_d': et_d, 'st_dt': st_dt, 'et_dt': et_dt}

    return res


def main():
    index = '399300'
    # stocks = get_stocks()
    # stocks = ['510880', '159915', '513100', '518880']
    # stocks = ['159901', '510880']
    # stocks = ['159901', '510880', '512510', '512220', '159938', '159909', '513500', '513880', '515770', '515350', '159810', '517350']
    # stocks = stocks[:200]
    # stocks = ['159909', '510170', '159915', '512220', '510880', '510300']
    #
    # stocks = [('159918', '159909', '512220', '159915', '510500', '510510', '159907', '159922', '510290', '159936',
    #            '159930', '159935', '159912', '159903', '159822', '513050', '159902', '510900', '512070', '510410')[19]]
    stocks = ['510900']
    # print(stocks)
    stocks.insert(0, index)
    run_backtrade(stocks)
    # s = get_data(stocks, freq='day', source='not db')
    # print(s)


if __name__ == '__main__':
    main()
